import cv2
import numpy as np
import matplotlib.pyplot as plt
import heapq
import math
import time

# ========================
# JPS 节点
# ========================
class Node:
    def __init__(self, x, y, g=0, h=0, parent=None):
        self.x = x
        self.y = y
        self.g = g
        self.h = h
        self.f = g + h
        self.parent = parent
    def __lt__(self, other):
        return self.f < other.f

# ========================
# JPS 算法
# ========================
class JPS:
    def __init__(self, grid):
        self.grid = grid
        self.rows, self.cols = grid.shape
        # 八方向 dx, dy, cost
        self.dirs = [
            (1,0,1),( -1,0,1),(0,1,1),(0,-1,1),
            (1,1,math.sqrt(2)),(-1,-1,math.sqrt(2)),
            (1,-1,math.sqrt(2)),(-1,1,math.sqrt(2))
        ]

    def in_bounds(self, x, y):
        return 0 <= x < self.cols and 0 <= y < self.rows

    def is_passable(self, x, y):
        return self.in_bounds(x, y) and self.grid[y, x] != 0

    def heuristic(self, x1, y1, x2, y2):
        dx, dy = abs(x1 - x2), abs(y1 - y2)
        return dx + dy + (math.sqrt(2) - 2) * min(dx, dy)

    # ========================
    # JPS 核心：沿方向跳点搜索
    # ========================
    def jump(self, x, y, dx, dy, end):
        nx, ny = x + dx, y + dy
        if not self.is_passable(nx, ny):
            return None
        if (nx, ny) == end:
            return (nx, ny)

        # 强制邻居（迫使分叉）
        if dx != 0 and dy != 0:
            # 对角线
            if (self.is_passable(nx - dx, ny + dy) and not self.is_passable(nx - dx, ny)) \
               or (self.is_passable(nx + dx, ny - dy) and not self.is_passable(nx, ny - dy)):
                return (nx, ny)
        else:
            # 直线
            if dx != 0:
                if (self.is_passable(nx + dx, ny + 1) and not self.is_passable(nx, ny + 1)) \
                   or (self.is_passable(nx + dx, ny - 1) and not self.is_passable(nx, ny - 1)):
                    return (nx, ny)
            else:
                if (self.is_passable(nx + 1, ny + dy) and not self.is_passable(nx + 1, ny)) \
                   or (self.is_passable(nx - 1, ny + dy) and not self.is_passable(nx - 1, ny)):
                    return (nx, ny)

        # 递归沿对角线跳
        if dx != 0 and dy != 0:
            if self.jump(nx, ny, dx, 0, end) is not None or self.jump(nx, ny, 0, dy, end) is not None:
                return (nx, ny)

        return self.jump(nx, ny, dx, dy, end)

    # 获取邻居跳点
    def identify_successors(self, node, end, g_score):
        successors = []
        for dx, dy, cost in self.dirs:
            result = self.jump(node.x, node.y, dx, dy, end)
            if result is not None:
                nx, ny = result
                tentative_g = g_score[(node.x, node.y)] + self.heuristic(node.x, node.y, nx, ny)
                successors.append(Node(nx, ny, tentative_g, self.heuristic(nx, ny, *end), node))
        return successors

    # ========================
    # 主搜索函数
    # ========================
    def find_path(self, start, end):
        open_list = []
        start_node = Node(*start, 0, self.heuristic(*start, *end))
        heapq.heappush(open_list, (start_node.f, start_node))
        g_score = {start: 0}

        while open_list:
            _, current = heapq.heappop(open_list)
            if (current.x, current.y) == end:
                path = []
                while current:
                    path.append((current.x, current.y))
                    current = current.parent
                return path[::-1]

            for neighbor in self.identify_successors(current, end, g_score):
                if (neighbor.x, neighbor.y) not in g_score or neighbor.g < g_score[(neighbor.x, neighbor.y)]:
                    g_score[(neighbor.x, neighbor.y)] = neighbor.g
                    heapq.heappush(open_list, (neighbor.f, neighbor))
        return []

# ========================
# 主程序
# ========================
# 读取地图
img = cv2.imread("map.png", cv2.IMREAD_GRAYSCALE)

# 定义可通行区域（去掉 0 和 192）
walkable_values = {255, 64, 32, 96, 120, 130, 140, 220}
grid = np.where(np.isin(img, list(walkable_values)), 1, 0)

# 膨胀障碍，考虑机器尺寸
robot_size = 19
obstacle = (1 - grid).astype(np.uint8)
kernel = np.ones((robot_size, robot_size), np.uint8)
inflated_obstacle = cv2.dilate(obstacle, kernel, iterations=1)
inflated_grid = 1 - inflated_obstacle

print("原始可通行格子数:", np.sum(grid))
print("膨胀后可通行格子数:", np.sum(inflated_grid))

# 初始化 JPS
jps = JPS(inflated_grid)

# ========================
# 点击交互
# ========================
start, end = None, None

def onclick(event):
    global start, end
    if event.xdata is None or event.ydata is None:
        return
    x, y = int(event.xdata), int(event.ydata)
    if inflated_grid[y, x] == 0:
        print("❌ 点击在障碍物上:", (x, y))
        return

    if start is None:
        start = (x, y)
        plt.scatter([x], [y], color="green", s=50)
        plt.draw()
        print("✅ 起点:", start)
    elif end is None:
        end = (x, y)
        plt.scatter([x], [y], color="blue", s=50)
        print("✅ 终点:", end)

        # 计算路径并测时间
        t0 = time.time()
        path = jps.find_path(start, end)
        t1 = time.time()
        elapsed = t1 - t0
        print(f"⏱ 路径计算耗时: {elapsed:.6f} 秒")

        if path:
            px, py = zip(*path)
            plt.plot(px, py, color="red")
            print("✅ 找到路径, 长度:", len(path))
        else:
            print("❌ 没有找到路径")
        plt.draw()

# ========================
# 可视化
# ========================
plt.imshow(inflated_grid, cmap="gray")
plt.title("点击选择起点 (绿色) 和终点 (蓝色)")
cid = plt.gcf().canvas.mpl_connect('button_press_event', onclick)
plt.show()
